// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
//

#include <kllm/kai/slots.h>

namespace kllm {

    void KaiSlots::save(const KaiRequest &req, KaiResponse&res) {
        std::string filename = req.slots_task().filename();
        if (!fs_validate_filename(filename)) {
            res.mutable_status()->set_code(ERROR_TYPE_INVALID_REQUEST);
            res.mutable_status()->set_errmsg("Invalid filename");
            return;
        }
        res.set_query_type(req.query_type());
        std::string filepath = req.slots_task().filepath();

        server_task task;
        task.type = SERVER_TASK_TYPE_SLOT_SAVE;
        *task.req.mutable_slots_task() = req.slots_task();

        const int id_task = _context->queue_tasks.post(task);
        _context->queue_results.add_waiting_task_id(id_task);

        TaskResult result = _context->queue_results.recv(id_task);
        _context->queue_results.remove_waiting_task_id(id_task);

        if (result.error) {
            res.mutable_status()->set_code(result.errcode);
            res.mutable_status()->set_errmsg(result.errmsg);
        } else {
            *res.mutable_slots_task() = result.slots_task;
        }
    }

    void KaiSlots::restore(const KaiRequest &req, KaiResponse &res) {
        std::string filename = req.slots_task().filename();
        if (!fs_validate_filename(filename)) {
            res.mutable_status()->set_code(ERROR_TYPE_INVALID_REQUEST);
            res.mutable_status()->set_errmsg("Invalid filename");
            return;
        }
        res.set_query_type(req.query_type());
        std::string filepath = req.slots_task().filepath();

        server_task task;
        task.type = SERVER_TASK_TYPE_SLOT_RESTORE;
        *task.req.mutable_slots_task() = req.slots_task();


        const int id_task = _context->queue_tasks.post(task);
        _context->queue_results.add_waiting_task_id(id_task);

        TaskResult result = _context->queue_results.recv(id_task);
        _context->queue_results.remove_waiting_task_id(id_task);

        if (result.error) {
            res.mutable_status()->set_code(result.errcode);
            res.mutable_status()->set_errmsg(result.errmsg);
        } else {
            *res.mutable_slots_task() = result.slots_task;
        }

    }

    void KaiSlots::erase(const KaiRequest &req, KaiResponse &res) {
        server_task task;
        task.type = SERVER_TASK_TYPE_SLOT_ERASE;
        *task.req.mutable_slots_task() = req.slots_task();
        res.set_query_type(req.query_type());
        const int id_task = _context->queue_tasks.post(task);
        _context->queue_results.add_waiting_task_id(id_task);

        TaskResult result = _context->queue_results.recv(id_task);
        _context->queue_results.remove_waiting_task_id(id_task);

        if (result.error) {
            res.mutable_status()->set_code(result.errcode);
            res.mutable_status()->set_errmsg(result.errmsg);
        } else {
            *res.mutable_slots_task() = result.slots_task;
        }
    }

}  // namespace kllm
